import sys

import numpy as np
from scipy.ndimage import uniform_filter

from agent import AGENT_PARAMS

sys.path.append('Regional_Interception_obstacles/')

from multi_target_prob_map import MultiTargetProbabilityMap


def situation_conversion(heatmap, threshold=0.8, neighborhood_size=3):
    """
    态势转换函数
    输入:
        heatmap: 敌方无人船热力态势矩阵 (2D numpy array)
        threshold: 敌方无人船出现的概率阈值，默认为0.8
        neighborhood_size: 邻域窗口的边长，默认为3
    输出:
        最可能敌方无人船位置的网格坐标 (row, col)
    """

    if not isinstance(heatmap, np.ndarray) or heatmap.ndim != 2:
        raise ValueError("heatmap必须是二维numpy数组！")

    # 1) 找到大于阈值的网格
    potential_mask = heatmap > threshold
    if not np.any(potential_mask):
        raise ValueError("没有任何网格的概率超过阈值，无法确定敌方位置。")

    # 2) 在整张热力图上做邻域平均
    # uniform_filter：在每个位置上做一个neighborhood_size的均值卷积
    local_mean = uniform_filter(heatmap, size=neighborhood_size, mode='constant')

    # 3) 只在潜在区域中寻找均值最高的网格
    local_mean_masked = np.where(potential_mask, local_mean, -np.inf)
    max_idx_flat = np.argmax(local_mean_masked)
    max_pos = np.unravel_index(max_idx_flat, heatmap.shape)

    return max_pos


def regional_interception_game_obstacles(enemy_pos1, enemy_velocity_ori,
                                         friend_pos, friend_speed, multi_target_map: MultiTargetProbabilityMap,
                                         scout_radius=30000, prediction_steps=100, step_size=1.0):
    """
    区域拦截博弈函数（避障版）- 性能优化版

    性能优化改进：
    1. 预计算重复使用的norm值，避免重复计算
    2. 预计算单位向量，消除重复计算
    3. 优化循环搜索，边计算边比较，减少内存占用
    4. 使用平方距离比较，避免不必要的开方运算
    5. 移除无效的函数调用

    输入：
        enemy_pos1: 敌方无人船当前位置(B)
        enemy_speed: 敌方无人船速度(矢量)
        obstacle_map: 不可通行坐标集合(np.ndarray)
        friend_pos: 我方无人船位置(A) 
        friend_speed: 我方无人船速度(标量)
        scout_radius: 其他敌方侦察范围半径
        prediction_steps: 预测轨迹的步数
        step_size: 每步时间间隔

    输出：
        apollonius_center: 阿氏圆圆心坐标（若存在）
        apollonius_radius: 阿氏圆半径（若存在）
        optimal_intercept_point: 最优拦截点（若存在）
        virtual_enemy_pos: 虚拟敌方位置（若存在）
        if_dangerous: 是否无解 (true--无解, false--有解)
    """
    # 预转换为numpy数组，保持数据类型一致性
    enemy_pos1 = np.array(enemy_pos1, dtype=float)
    friend_pos = np.array(friend_pos, dtype=float)
    enemy_velocity = np.array(enemy_velocity_ori, dtype=float)

    # 预计算基础向量和距离
    enemy_to_friend_vec = friend_pos - enemy_pos1
    distance_enemy_friend = np.linalg.norm(enemy_to_friend_vec)

    LOCK_RADIUS = scout_radius
    DESIRED_DISTANCE = 35000

    # 优化1：预计算单位向量，避免重复计算
    unit_vec = enemy_to_friend_vec / distance_enemy_friend

    # 情况1：敌我距离小于锁定半径
    if distance_enemy_friend <= LOCK_RADIUS:
        if distance_enemy_friend < DESIRED_DISTANCE:
            print("敌我之间距离小于35km，无法保持安全距离")
        optimal_intercept_point = enemy_pos1 + unit_vec * DESIRED_DISTANCE
        return None, None, optimal_intercept_point, None, False

    # 计算虚拟敌方位置
    virtual_enemy_pos = enemy_pos1 + unit_vec * DESIRED_DISTANCE

    # 优化2：预计算enemy_velocity的norm，避免重复计算3次
    enemy_velocity_norm = np.linalg.norm(enemy_velocity)

    # 情况2：敌方速度为零
    if enemy_velocity_norm == 0:
        optimal_intercept_point = enemy_pos1 + unit_vec * DESIRED_DISTANCE

        # 检查optimal_intercept_point是否为NaN
        if np.isnan(optimal_intercept_point).any():
            return None, None, None, optimal_intercept_point, True

        if multi_target_map.can_pass(optimal_intercept_point):
            return None, None, optimal_intercept_point, virtual_enemy_pos, False
        else:
            return None, None, None, optimal_intercept_point, True

    # 使用预计算的norm值
    enemy_dir_unit = enemy_velocity / enemy_velocity_norm
    k = friend_speed / enemy_velocity_norm

    # 计算阿波罗尼奥斯圆参数
    apollonius_center = (virtual_enemy_pos - k ** 2 * friend_pos) / (1 - k ** 2)
    apollonius_radius = (k * np.linalg.norm(virtual_enemy_pos - friend_pos)) / abs(1 - k ** 2)

    # 求解阿氏圆与敌方方向直线的交点
    dp = virtual_enemy_pos - apollonius_center
    a = np.dot(enemy_dir_unit, enemy_dir_unit)  # 实际等于1.0
    b = 2 * np.dot(dp, enemy_dir_unit)
    c = np.dot(dp, dp) - apollonius_radius ** 2

    discriminant = b ** 2 - 4 * a * c
    if discriminant < 0:
        print("阿氏圆与敌方方向直线无交点，无法拦截！")
        return apollonius_center, apollonius_radius, None, virtual_enemy_pos, True

    sqrt_disc = np.sqrt(discriminant)
    t1 = (-b + sqrt_disc) / (2 * a)
    t2 = (-b - sqrt_disc) / (2 * a)

    intercept1 = virtual_enemy_pos + t1 * enemy_dir_unit
    intercept2 = virtual_enemy_pos + t2 * enemy_dir_unit

    # 优化3：使用平方距离比较，避免开方运算
    dist1_sq = np.sum((intercept1 - virtual_enemy_pos) ** 2)
    dist2_sq = np.sum((intercept2 - virtual_enemy_pos) ** 2)
    candidate_intercept = intercept1 if dist1_sq < dist2_sq else intercept2

    # 优化4：移除无用的get_obstacle_map()调用（原代码未使用返回值）
    # multi_target_map.get_obstacle_map()  # 已删除

    # 检查optimal_intercept_point是否为NaN
    if np.isnan(candidate_intercept).any():
        return apollonius_center, apollonius_radius, candidate_intercept, virtual_enemy_pos, False

    if multi_target_map.can_pass(candidate_intercept):
        return apollonius_center, apollonius_radius, candidate_intercept, virtual_enemy_pos, False

    # 优化5：次优搜索 - 边计算边比较，避免存储所有候选点
    best_point = None
    min_time_diff = float('inf')
    found_valid_point = False

    for i in range(1, prediction_steps + 1):
        t = i * step_size
        predicted_point = virtual_enemy_pos + enemy_velocity * t

        if np.isnan(predicted_point).any():
            continue

        if not multi_target_map.can_pass(predicted_point):
            continue

        found_valid_point = True
        dist_friend = np.linalg.norm(predicted_point - friend_pos)
        friend_time = dist_friend / friend_speed
        time_diff = abs(friend_time - t)

        # 直接比较，不存储所有候选点
        if time_diff < min_time_diff:
            min_time_diff = time_diff
            best_point = predicted_point

    if not found_valid_point:
        print("所有可预测点都在侦察范围内，无法拦截")
        return apollonius_center, apollonius_radius, candidate_intercept, virtual_enemy_pos, True

    return apollonius_center, apollonius_radius, best_point, virtual_enemy_pos, False


def regional_interception_game(enemy_pos, enemy_velocity_ori,
                               friend_pos, friend_speed = 2 * AGENT_PARAMS['USV_MAX_VEL'],
                               scout_radius=AGENT_PARAMS['BSV_SCAN_RANGE'] + 1000):
    """
    简化版区域拦截博弈函数（无避障功能）

    输入：
        enemy_pos: 敌方无人船当前位置(B)
        enemy_velocity_ori: 敌方无人船速度(矢量)
        friend_pos: 我方无人船位置(A)
        friend_speed: 我方无人船速度(标量)
        scout_radius: 侦察范围半径
        prediction_steps: 预测轨迹的步数
        step_size: 每步时间间隔

    输出：
        apollonius_center: 阿氏圆圆心坐标
        apollonius_radius: 阿氏圆半径
        optimal_intercept_point: 最优拦截点
        virtual_enemy_pos: 虚拟敌方位置
        if_dangerous: 是否无解 (True-无解, False-有解)
    """
    # 预转换为numpy数组
    enemy_pos = np.array(enemy_pos, dtype=float)
    friend_pos = np.array(friend_pos, dtype=float)
    enemy_velocity = np.array(enemy_velocity_ori, dtype=float)

    # 计算基础向量和距离
    enemy_to_friend_vec = friend_pos - enemy_pos
    distance_enemy_friend = np.linalg.norm(enemy_to_friend_vec)

    LOCK_RADIUS = scout_radius
    DESIRED_DISTANCE = 35000

    # 计算单位向量
    if distance_enemy_friend > 0:
        unit_vec = enemy_to_friend_vec / distance_enemy_friend
    else:
        unit_vec = np.array([1, 0])  # 默认方向

    # 情况1：敌我距离小于锁定半径
    if distance_enemy_friend <= LOCK_RADIUS:
        optimal_intercept_point = enemy_pos + unit_vec * DESIRED_DISTANCE
        return None, None, optimal_intercept_point, None, False

    # 计算虚拟敌方位置
    virtual_enemy_pos = enemy_pos + unit_vec * DESIRED_DISTANCE

    # 计算敌方速度大小
    enemy_velocity_norm = np.linalg.norm(enemy_velocity)

    # 情况2：敌方速度为零
    if enemy_velocity_norm == 0:
        optimal_intercept_point = virtual_enemy_pos
        return None, None, optimal_intercept_point, virtual_enemy_pos, False

    # 计算速度比
    enemy_dir_unit = enemy_velocity / enemy_velocity_norm
    k = friend_speed / enemy_velocity_norm

    # 计算阿波罗尼奥斯圆参数
    apollonius_center = (virtual_enemy_pos - k ** 2 * friend_pos) / (1 - k ** 2)
    apollonius_radius = (k * np.linalg.norm(virtual_enemy_pos - friend_pos)) / abs(1 - k ** 2)

    # 求解阿氏圆与敌方方向直线的交点
    dp = virtual_enemy_pos - apollonius_center
    a = 1.0  # 因为enemy_dir_unit是单位向量
    b = 2 * np.dot(dp, enemy_dir_unit)
    c = np.dot(dp, dp) - apollonius_radius ** 2

    discriminant = b ** 2 - 4 * a * c
    if discriminant < 0:
        return apollonius_center, apollonius_radius, None, virtual_enemy_pos, True

    sqrt_disc = np.sqrt(discriminant)
    t1 = (-b + sqrt_disc) / (2 * a)
    t2 = (-b - sqrt_disc) / (2 * a)

    intercept1 = virtual_enemy_pos + t1 * enemy_dir_unit
    intercept2 = virtual_enemy_pos + t2 * enemy_dir_unit

    # 选择距离虚拟敌位置更近的交点
    dist1_sq = np.sum((intercept1 - virtual_enemy_pos) ** 2)
    dist2_sq = np.sum((intercept2 - virtual_enemy_pos) ** 2)
    optimal_intercept_point = intercept1 if dist1_sq < dist2_sq else intercept2

    return apollonius_center, apollonius_radius, optimal_intercept_point, virtual_enemy_pos, False


import math


def calculate_interception_point(
        enemy_pos: list[float],
        enemy_vel: list[float],
        friend_pos: list[float],
        v_max: float = AGENT_PARAMS['USV_MAX_VEL'],
        lock_range: float = AGENT_PARAMS['USV_LOCK_RANGE'] - 5000,
        enemy_scan_range: float = AGENT_PARAMS['BSV_SCAN_RANGE'] + 3500,
        alpha: float = None,  # None表示自动计算alpha
        defense_line_x: float = 50000.0
):
    """
    计算拦截点 (增强鲁棒性版本)
    :param friend_pos: 我方位置 (x0, y0)
    :param enemy_pos: 敌方位置 (x1, y1)
    :param v_max: 我方最大速度（标量）
    :param enemy_vel: 敌方速度矢量 (vx, vy)
    :param lock_range: 我方射程
    :param enemy_scan_range: 敌方射程
    :param alpha: 预测权重系数 (None=自动计算, 0=纯保守策略, 1=纯预测策略)
    :return: 拦截点坐标 (x, y)
    """
    # 计算敌我位置矢量差
    dx = enemy_pos[0] - friend_pos[0]
    dy = enemy_pos[1] - friend_pos[1]
    dist0 = math.sqrt(dx ** 2 + dy ** 2)

    # 预先计算保守拦截点 (敌我连线延长线方向)
    reverse_ux = (friend_pos[0] - enemy_pos[0]) / dist0
    reverse_uy = (friend_pos[1] - enemy_pos[1]) / dist0
    conservative_x = enemy_pos[0] + lock_range * reverse_ux
    conservative_y = enemy_pos[1] + lock_range * reverse_uy

    # 如果初始距离已在侦察范围内，直接返回保守拦截点
    if dist0 <= enemy_scan_range:
        return [conservative_x, conservative_y], True

    # 自动计算alpha值 (150000m->0.4, 40000m->0.8)
    if alpha is None:
        # 距离范围定义
        MAX_DIST = 150000.0  # 150公里
        MIN_DIST = 40000.0  # 40公里

        # 将距离映射到[0,1]区间并反转 (距离越大，值越小)
        normalized_dist = (dist0 - MIN_DIST) / (MAX_DIST - MIN_DIST)
        normalized_dist = max(0.0, min(1.0, normalized_dist))  # 夹紧到[0,1]
        reversed_dist = 1.0 - normalized_dist

        # 计算alpha值 (0.4到0.8之间)
        alpha_val = 0.4 + 0.4 * reversed_dist
    else:
        alpha_val = alpha

    # 计算单位方向矢量 (从P0指向P1)
    ux = dx / dist0
    uy = dy / dist0

    # 计算敌方速度在单位方向矢量上的投影 (平行分量)
    enemy_vel_para = enemy_vel[0] * ux + enemy_vel[1] * uy

    # 计算敌方速度的垂直分量大小
    enemy_vel_perp_x = enemy_vel[0] - enemy_vel_para * ux
    enemy_vel_perp_y = enemy_vel[1] - enemy_vel_para * uy
    enemy_vel_perp_sq = enemy_vel_perp_x ** 2 + enemy_vel_perp_y ** 2

    # 检查是否无法拦截
    B = v_max ** 2 - enemy_vel_perp_sq
    if B < 0 or (B >= 0 and math.sqrt(B) - enemy_vel_para < 0.05):
        # 无法拦截时，计算敌方速度方向与防御线的交点
        enemy_speed = math.sqrt(enemy_vel[0] ** 2 + enemy_vel[1] ** 2)

        if enemy_speed < 1e-5:  # 敌方静止
            # 返回防御线上与敌方当前位置y相同的点
            return [defense_line_x, enemy_pos[1]], False
        else:
            # 计算敌方速度方向单位向量
            enemy_dir_x = enemy_vel[0] / enemy_speed
            enemy_dir_y = enemy_vel[1] / enemy_speed

            # 计算与防御线交点
            t = (defense_line_x - enemy_pos[0]) / enemy_dir_x if enemy_dir_x != 0 else float('inf')

            # 确保t是正值（敌方朝防御线方向移动）
            if t > 0:
                y_intersect = enemy_pos[1] + t * enemy_dir_y
                return [alpha_val * defense_line_x + (1 - alpha_val) * conservative_x,
                        alpha_val * y_intersect + (1 - alpha_val) * conservative_y], False
            else:
                # 如果t为负值，说明敌方背向防御线移动，返回防御线上与敌方当前位置y相同的点
                return [friend_pos[0], enemy_pos[1]], False

    # 计算最大相对速度分量
    a_max = math.sqrt(B) - enemy_vel_para

    # 计算拦截所需时间
    t = (dist0 - lock_range) / a_max

    # 计算预测拦截点（敌方在t时刻的位置）
    predicted_x = enemy_pos[0] + enemy_vel[0] * t
    predicted_y = enemy_pos[1] + enemy_vel[1] * t

    # 加权平均：混合预测点和保守点
    intercept_x = alpha_val * predicted_x + (1 - alpha_val) * conservative_x
    intercept_y = alpha_val * predicted_y + (1 - alpha_val) * conservative_y

    return [intercept_x, intercept_y], True
